"""结合提示词和标题问题，对文本进行分类"""
import os

from dotenv import load_dotenv
from openai import OpenAI

load_dotenv()

SYSTEM_PROMPT = """
你是一个专业的中文新闻标题分类器。你的任务是根据用户输入的中文新闻标题，将其准确分类到以下10个类别之一：
finance（金融）、realty（房地产）、stocks（股票）、education（教育）、science（科学）、society（社会）、politics（政治）、sports（体育）、game（游戏）、entertainment（娱乐）。请注意：
只输出类别名称的英文单词（例如：finance、sports等），不要输出任何其他文本、解释或标点符号。
分类必须严格基于标题内容，并仅限于上述类别。
如果标题不明显匹配任何类别，基于内容选择最接近的类别。
现在，请提供新闻标题，我将返回对应的类别。
"""

def deepseek_predict(prompt: str):
    api_key = os.getenv("DEEPSEEK_API_KEY")
    base_url = os.getenv("DEEPSEEK_BASE_URL")
    client = OpenAI(api_key=api_key, base_url=base_url)

    response = client.chat.completions.create(
        model="deepseek-chat",
        messages=[
            {"role": "system", "content": "You are a helpful assistant"},
            {"role": "user", "content": prompt},
        ],
        stream=False
    )
    return response.choices[0].message.content

if __name__ == '__main__':
    print(deepseek_predict(SYSTEM_PROMPT + "新闻标题：中国股市再次创新高"))
